#!/usr/bin/env python3
"""
Command-line script for importing OpenCTI data into the Knowledge Graph.
"""
import os
import sys
import argparse
import logging
from datetime import datetime
import time
from dotenv import load_dotenv

from connector import OpenCTIConnector

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler('opencti_import.log')
    ]
)
logger = logging.getLogger(__name__)

def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description='Import data from OpenCTI to the Knowledge Graph')
    
    # Mode selection
    parser.add_argument('--mode', choices=['once', 'daemon'], default='once',
                        help='Run mode: once (single import) or daemon (continuous imports)')
    
    # OpenCTI connection settings
    parser.add_argument('--opencti-url', help='OpenCTI API URL')
    parser.add_argument('--opencti-token', help='OpenCTI API token')
    
    # Neo4j connection settings
    parser.add_argument('--neo4j-uri', help='Neo4j database URI')
    parser.add_argument('--neo4j-user', help='Neo4j username')
    parser.add_argument('--neo4j-password', help='Neo4j password')
    
    # Import options
    parser.add_argument('--interval', type=int, default=24,
                       help='Interval between imports in hours (daemon mode only)')
    parser.add_argument('--limit', type=int, default=500,
                       help='Maximum number of objects to import per type')
    parser.add_argument('--types', nargs='+', 
                       help='STIX types to import (default: all supported types)')
    
    return parser.parse_args()

def load_config(args):
    """Load configuration from environment variables and command-line arguments."""
    # Load from .env file if present
    load_dotenv()
    
    # Configuration dictionary with priorities: args > env > defaults
    config = {
        # Neo4j settings
        'neo4j_uri': args.neo4j_uri or os.getenv('NEO4J_URI') or 'bolt://localhost:7687',
        'neo4j_user': args.neo4j_user or os.getenv('NEO4J_USER') or 'neo4j',
        'neo4j_password': args.neo4j_password or os.getenv('NEO4J_PASSWORD') or 'password',
        
        # OpenCTI settings
        'opencti_url': args.opencti_url or os.getenv('OPENCTI_URL'),
        'opencti_token': args.opencti_token or os.getenv('OPENCTI_TOKEN'),
        
        # Import settings
        'mode': args.mode,
        'interval': args.interval,
        'limit': args.limit,
        'types': args.types
    }
    
    # Validate required settings
    missing = []
    if not config['opencti_url']:
        missing.append('OpenCTI URL')
    if not config['opencti_token']:
        missing.append('OpenCTI API token')
    
    if missing:
        logger.error(f"Missing required configuration: {', '.join(missing)}")
        logger.error("Please provide via command-line arguments or environment variables")
        sys.exit(1)
    
    return config

def run_import(config):
    """Run the import process."""
    try:
        # Initialize connector
        connector = OpenCTIConnector(
            neo4j_uri=config['neo4j_uri'],
            neo4j_user=config['neo4j_user'],
            neo4j_password=config['neo4j_password'],
            opencti_url=config['opencti_url'],
            opencti_token=config['opencti_token']
        )
        
        if config['mode'] == 'daemon':
            logger.info(f"Starting OpenCTI import daemon (interval: {config['interval']} hours)")
            connector.run_incremental_import(
                interval_hours=config['interval'],
                types=config['types']
            )
        else:
            # Run a single import
            logger.info("Starting one-time OpenCTI import")
            start_time = time.time()
            stats = connector.import_stix_objects(
                limit=config['limit'],
                types=config['types']
            )
            duration = time.time() - start_time
            
            # Output results
            logger.info(f"Import completed in {duration:.2f} seconds")
            logger.info(f"Imported {stats['total_imported']} objects")
            for type_name, count in stats.get('by_type', {}).items():
                logger.info(f"  - {type_name}: {count} objects")
            
            # Clean up
            connector.close()
            
    except Exception as e:
        logger.error(f"Import failed: {str(e)}")
        sys.exit(1)

def main():
    """Main entry point."""
    args = parse_args()
    config = load_config(args)
    run_import(config)

if __name__ == "__main__":
    main() 